# data_reranker.py
# usage: 根据传递过来的documents和question，调用reranker模型排序
# 运行前，先执行./src/query/query_vector.py 生成tmp_filtered_docs.txt
# python ./src/query/query_reranker.py --query_text "请问在2025年的农林牧渔专题中，D区办结工单量相较于A区多出多少（若A区更多则请用负数表示）？"


import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
from config import project_config
import argparse
import json
import requests
from utils import debug_print


def query_reranker(request_body, debug=None):

    debug_print(debug,f"✅ 加载请求数据成功：{request_body}")
    # print(f"查询: {request_body['query']}\n")

    # 2. 准备请求参数
    url = "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
    headers = {
        "Content-Type": "application/json",
        # "Authorization": "Bearer sk-..."  # 替换为你的实际 api_key
        "Authorization": "Bearer "+project_config.llm_api_key
    } 

    # 3. 发送 POST 请求
    try:
        response = requests.post(url, headers=headers, json=request_body, timeout=30)
        response.raise_for_status()  # 检查 HTTP 错误
    except requests.exceptions.RequestException as e:
        print(f"❌ 请求失败：{e}")
        return None

    # 4. 解析响应
    try:
        result = response.json()
    except json.JSONDecodeError:
        print("❌ 响应不是有效的 JSON 格式：")
        print(response.text)
        return None

    debug_print(debug,"✅ Reranker API 调用成功！\n")

    # 5. 提取 output.results  
    results = result.get("output", {}).get("results", [])
    if not results:
        debug_print(debug,"⚠️  返回结果中没有 'output.results' 字段或为空。")
        return None

    debug_print(debug,f"rerank result: {results}")


    best_result = max(results, key=lambda x: x["relevance_score"])
    best_index = best_result["index"]
    best_score = best_result["relevance_score"]

    documents = request_body["input"]["documents"]
    if best_index >= len(documents):
        print(f"❌ 索引越界：best_index={best_index}, 文档数量={len(documents)}")
        return None

    # best_document = documents[best_index]
    best_document_str = documents[best_index]  # 这是一个 JSON 字符串
    best_document = json.loads(best_document_str)  # 转成字典
    # 6. 输出最相关文档
    debug_print(debug,"=" * 60)
    debug_print(debug,"🏆 最相关文档（得分最高）:")
    debug_print(debug,f"Index: {best_index}")
    debug_print(debug,f"Relevance Score: {best_score:.6f}")
    debug_print(debug,"-" * 60)
    debug_print(debug,best_document["DDL"])
    debug_print(debug,"=" * 60)

    return best_document["DDL"]


def main():

    # 创建参数解析器
    parser = argparse.ArgumentParser(description="检索向量库中相似度最高的内容。")
    parser.add_argument('--query_text', type=str, required=True, help='需要检索的文本字符串。' )
    parser.add_argument('--debug', action='store_true', help='启用调试模式，显示详细日志。')
    args = parser.parse_args()
    text_to_query = args.query_text
    
    input_file = 'tmp_filtered_docs.txt'
    
    # 1. 读取输入文件
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            docs = json.load(f)
    except FileNotFoundError:
        print(f"错误：文件 {input_file} 未找到！")
    except json.JSONDecodeError as e:
        print(f"错误：文件格式不是合法的 JSON。{e}")

    request_body = {
        "model": project_config.reranker_model,
        "input": {
            "query": text_to_query,
            "documents": docs
        }
    }
    final_doc = query_reranker(request_body, args.debug)
    print(f"final_doc={final_doc}")

# 主程序入口
if __name__ == "__main__":
    main()
